package com.aispeech.asr.clib.post;

import com.sun.jna.Native;
import lombok.extern.slf4j.Slf4j;

import java.util.List;

@Slf4j
public class Post {
    private static IPost iPost;
    private static long ptPostCfgDefault16k;
    private static long ptPostCfgDefault8k;
    private static long ptPostCfgEn;

    private long ptPost;

    public Post() {
        this(null, 16000);
    }

    public Post(boolean useAux) {
        this(null, 16000, useAux);
    }

    public Post(String lang, int sampleRate) {
        this(lang, sampleRate, false);
    }

    public Post(String lang, int sampleRate, boolean useAux) {
        ptPost = iPost.lasr_post_process_engine_new(getPostCfgByLang(lang, sampleRate), useAux ? 1 : 0);
        assert ptPost != 0;
    }

    public static void initialize(String resPath, String ipcPath) {
        initialize(resPath, ipcPath, false);
    }

    public static void initialize(String resPath, String ipcPath, boolean initAsrpp) {
        iPost = Native.load("post", IPost.class);

        if (iPost == null) {
            throw new RuntimeException("Load Post library failed");
        }

        ptPostCfgDefault16k = iPost.lasr_post_process_static_dependency_init(
                String.format("%s/cn/txtpost.cfg", resPath),
                initAsrpp ? String.format("%s/cn/asrpp_16000.cfg", resPath) : null,
                ipcPath);
        if (ptPostCfgDefault16k <= 0) {
            throw new RuntimeException("Initialize Post cn library failed");
        }

        ptPostCfgDefault8k = iPost.lasr_post_process_static_dependency_init(
                String.format("%s/cn/txtpost.cfg", resPath),
                initAsrpp ? String.format("%s/cn/asrpp_8000.cfg", resPath) : null,
                ipcPath);
        if (ptPostCfgDefault8k <= 0) {
            throw new RuntimeException("Initialize Post cn 8k library failed");
        }

        ptPostCfgEn = iPost.lasr_post_process_static_dependency_init(String.format("%s/en/txtpost.cfg", resPath), null, ipcPath);
        if (ptPostCfgEn <= 0) {
            throw new RuntimeException("Initialize Post en library failed");
        }

        log.info("Post library loaded.");
    }

    private static long getPostCfgByLang(String lang, int sampleRate) {
        if ((lang != null && "en".equals(lang))) {
            return ptPostCfgEn;
        } else {
            if (sampleRate == 8000) {
                return ptPostCfgDefault8k;
            } else {
                return ptPostCfgDefault16k;
            }
        }
    }

    public static PostResult process(AlignInfo alignInfo, PostEnv env) {
        return process(alignInfo, env, null);
    }

    public static PostResult process(AlignInfo alignInfo, PostEnv env, String lang) {
//        log.warn(">>>>> begin new post...");
        Post post = new Post(lang, 16000, false); //兼容一句话的使用，不需要asrpp的特性。
//        log.warn(">>>>> begin process post...");
        PostResult result = process(post, alignInfo, env);
//        log.warn(">>>>> begin delete post...");
        post.Delete();
        return result;
    }

    public static PostResult process(Post post, AlignInfo alignInfo, PostEnv env) {
        List<AlignElement> alignElements = alignInfo.getAlignElements();
        PostResult postResult = new PostResult();

        int size = alignElements.size();

        if (size <= 0) {
            return postResult;
        }

        long ptPost = post.ptPost;

        final AlignInfoStruct.ByReference alignInfoStruct = new AlignInfoStruct.ByReference();
        alignInfoStruct.aligns = new AlignElementStruct.ByReference();
        AlignElementStruct[] alignElementStructs = (AlignElementStruct[]) alignInfoStruct.aligns.toArray(size);
        alignInfoStruct.len = size;
        for (int i = 0; i < size; i++) {
            AlignElement element = alignElements.get(i);

            alignElementStructs[i].timeStart = element.begin;
            alignElementStructs[i].timeEnd = element.end;
            alignElementStructs[i].uid = element.speaker;
            alignElementStructs[i].words = element.rec;
        }

        String sEnv = env.toString();
//        log.info("post env: {}", sEnv);

//        log.warn(">>>>> pre lasr_post_process_engine_feed_align...");
        PostResultStruct.ByReference postResultStruct = iPost.lasr_post_process_engine_feed_align(
                ptPost,
                alignInfoStruct,
                sEnv
        );
//        log.warn(">>>>> after lasr_post_process_engine_feed_align...");

        postResult.setStatus(postResultStruct.error_code);
        if (postResult.getStatus() == 0)
            generatePostResult(postResult, postResultStruct);

        return postResult;
    }

    private static void generatePostResult(PostResult postResult, PostResultStruct postResultStruct) {
        AlignOutputStruct segmentResult = postResultStruct.segment_result;

        if (segmentResult != null && segmentResult.aligns != null && segmentResult.len > 0) {
            AlignElementStruct[] segmentElementStructs = (AlignElementStruct[]) postResultStruct
                    .segment_result.aligns.toArray(postResultStruct.segment_result.len);

            for (AlignElementStruct segmentElementStruct : segmentElementStructs) {
                postResult.addSegmentElement(segmentElementStruct.timeStart, segmentElementStruct.timeEnd,
                        segmentElementStruct.uid, segmentElementStruct.words, segmentElementStruct.pinyin, segmentElementStruct.pos);
            }
        } else {
            //log.warn("[ generatePostResult ] segment_result: {}", postResultStruct.segment_result);
        }

        AlignOutputStruct mergedResult = postResultStruct.merge_result;
        if (mergedResult != null && mergedResult.aligns != null && mergedResult.len > 0) {
            AlignElementStruct[] mergedElementStructs = (AlignElementStruct[]) postResultStruct
                    .merge_result.aligns.toArray(postResultStruct.merge_result.len);

            for (AlignElementStruct mergedElementStruct : mergedElementStructs) {
                postResult.addMergedElement(mergedElementStruct.timeStart, mergedElementStruct.timeEnd,
                        mergedElementStruct.uid, mergedElementStruct.words, mergedElementStruct.pinyin, mergedElementStruct.pos);
            }
        } else {
            log.info("[ generatePostResult ] merge_result is null");
        }

        postResult.addMetrics("duration_txtsmooth", postResultStruct.post_proc_detail.txtsmooth_time);
        postResult.addMetrics("duration_punc", postResultStruct.post_proc_detail.punc_time);
        postResult.addMetrics("duration_smooth_punc", postResultStruct.post_proc_detail.smooth_punc_unite_time);
        postResult.addMetrics("duration_itn", postResultStruct.post_proc_detail.itn_time);
        postResult.addMetrics("duration_pinyin", postResultStruct.post_proc_detail.pinyin_time);
        postResult.addMetrics("duration_hotwords", postResultStruct.post_proc_detail.hotwords_time);
        postResult.addMetrics("duration_sensitive", postResultStruct.post_proc_detail.sensitive_time);
        postResult.addMetrics("duration_paragraph", postResultStruct.post_proc_detail.paragraph_time);

        postResult.addMetrics("new_start_ts", postResultStruct.post_proc_detail.new_start_ts);
        postResult.addMetrics("feed_start_ts", postResultStruct.post_proc_detail.feed_start_ts);
        postResult.addMetrics("smooth_punc_start_ts", postResultStruct.post_proc_detail.smooth_punc_start_ts);
        postResult.addMetrics("smooth_start_ts", postResultStruct.post_proc_detail.smooth_start_ts);
        postResult.addMetrics("punc_start_ts", postResultStruct.post_proc_detail.punc_start_ts);
        postResult.addMetrics("pinyin_start_ts", postResultStruct.post_proc_detail.pinyin_start_ts);
        postResult.addMetrics("hotword_start_ts", postResultStruct.post_proc_detail.hotword_start_ts);
        postResult.addMetrics("paragraph_start_ts", postResultStruct.post_proc_detail.paragraph_start_ts);
        postResult.addMetrics("itn_start_ts", postResultStruct.post_proc_detail.itn_start_ts);
        postResult.addMetrics("sensitive_start_ts", postResultStruct.post_proc_detail.sensitive_start_ts);
    }

    public static void destroy() {
        if (ptPostCfgDefault16k > 0)
            iPost.lasr_post_process_static_dependency_destroy(ptPostCfgDefault16k);
        if (ptPostCfgDefault8k > 0)
            iPost.lasr_post_process_static_dependency_destroy(ptPostCfgDefault8k);
        if (ptPostCfgEn > 0)
            iPost.lasr_post_process_static_dependency_destroy(ptPostCfgEn);
    }

    public PostResult post(AlignInfo alignInfo, PostEnv env) {
        return process(this, alignInfo, env);
    }

    public void Delete() {
        if (ptPost != 0) {
            iPost.lasr_post_process_engine_delete(ptPost);
        }
    }

    public AsrppResult asrpp(byte[] audioBuffer) {
        AsrppResultStruct asrppResultStruct = iPost.lasr_asrpp_process_engine_feed_align(ptPost, audioBuffer, audioBuffer.length);
        AsrppResult result = new AsrppResult(asrppResultStruct.error_code, asrppResultStruct.asrpp_time, asrppResultStruct.asrpp_feed_start_ts);

        AsrppInfoStruct asrppInfoStruct = asrppResultStruct.asrpp_result;
        if (asrppInfoStruct != null && asrppInfoStruct.attribs != null && asrppInfoStruct.len > 0) {
            AsrppElementStruct[] asrppElementStructs = (AsrppElementStruct[]) asrppInfoStruct.attribs.toArray(asrppInfoStruct.len);

            for (AsrppElementStruct elementStruct : asrppElementStructs) {
                result.addElement(elementStruct.emotion, elementStruct.age, elementStruct.gender);
            }
        } else {
            log.info("[ asrpp ] len: {}, result: {}", asrppInfoStruct.len, asrppInfoStruct.attribs);
        }

        return result;
    }
}
